# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import numpy as np
from torch.autograd import Variable
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pylab as plt
# 1.建立数据集
x = torch.unsqueeze(torch.linspace(-1,1,100), dim = 1)
y = pow(x, 2) + 0.2*torch.rand(x.size())
x,y = Variable(x), Variable(y)
# plt.scatter(x.data.numpy(), y.data.numpy())
# plt.show()
# 2.建立神经网络

class Model(nn.Module):
    def __init__(self, n_feature, n_hidden, n_output):
        super(Model, self).__init__()
        self.hidden = torch.nn.Linear(n_feature, n_hidden)
        self.predict = torch.nn.Linear(n_hidden, n_output)
    def forward(self, x):
        #正向传播输入值，神经网络分析出输出值/
        x = F.relu(self.hidden(x))                #激励函数
        x = self.predict(x)
        return x
model = Model(n_feature=1, n_hidden=10, n_output=1)
# 3.训练网络与可视化
# optimizer是训练的工具
optimizer = optim.SGD(model.parameters(), lr=0.5)  #传入所有参数与学习率
loss_func = torch.nn.MSELoss()                     #预测值和真实值的误差计算公式（均方差）
plt.ion()
plt.show()
for i in range(100):
    prediction = model(x)             # 喂给 model 训练数据 x, 输出预测值
    loss = loss_func(prediction, y)   #计算两者误差
    optimizer.zero_grad()
    loss.backward()                   #误差反向传播，计算参数更新值
    optimizer.step()                  #将参数更新的值施加到model的parameters上
    if i % 5 == 0:
        plt.cla()
        plt.scatter(x.data.numpy(), y.data.numpy())
        plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
        # plt.text(0.5, 0, 'Loss=%.4f' % loss.data[0], fontdict={'size': 20, 'color':  'red'})
        plt.pause(0.1)



'''
x_train = np.array([[3.3], [4.4], [5.5], [6.71], [6.93], [4.168],
                    [9.779], [6.182], [7.59], [2.167], [7.042],
                    [10.791], [5.313], [7.997], [3.1]], dtype=np.float32)
x_train = torch.from_numpy(x_train)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1,20,5)
        self.conv2 = nn.Conv2d(20,20,5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))
class Test(nn.Module):
    def __init__(self):
        super(Test, self).__init__()
        self.add_module("conv", nn.Conv2d(10,20,4))
model = Model()
print(model)
test = Test()
print(test)
'''